﻿Imports System
Imports System.Collections
Imports System.Collections.Generic
Imports System.IO
Imports System.Workflow.ComponentModel

Namespace System.Workflow.Runtime
    Friend NotInheritable Class WorkflowStateRollbackService
        ' Methods
        Public Sub New(ByVal workflowExecutor As WorkflowExecutor)
            Me.workflowExecutor = workflowExecutor
        End Sub

        Friend Sub CheckpointInstanceState()
            Me.clonedInstanceStateStream = New MemoryStream(&H2800)
            Me.workflowExecutor.RootActivity.Save(Me.clonedInstanceStateStream)
            Me.workflowDefinition = Me.workflowExecutor.WorkflowDefinition
            Me.completedContextActivities = DirectCast(Me.workflowExecutor.CompletedContextActivities.Clone, Hashtable)
            Me.clonedInstanceStateStream.Position = 0
        End Sub

        Friend Sub DisposeCheckpointState()
            Me.clonedInstanceStateStream = Nothing
        End Sub

        Friend Sub RequestRevertToCheckpointState(ByVal currentActivity As Activity, ByVal callbackHandler As EventHandler([Of] EventArgs), ByVal callbackData As EventArgs, ByVal suspendOnRevert As Boolean, ByVal suspendInfo As String)
            If (Me.clonedInstanceStateStream Is Nothing) Then
                Throw New InvalidOperationException(ExecutionStringManager.InvalidRevertRequest)
            End If
            Me.activityContextId = ContextActivityUtils.ContextId(ContextActivityUtils.ContextActivity(currentActivity))
            Me.activityQualifiedName = currentActivity.QualifiedName
            Me.callbackData = callbackData
            Me.callbackHandler = callbackHandler
            Me.suspendOnRevert = suspendOnRevert
            Me.suspendOnRevertInfo = suspendInfo
            Me.isInstanceStateRevertRequested = True
            Me.workflowExecutor.Scheduler.CanRun = False
        End Sub

        Friend Sub RevertToCheckpointState()
            Dim rootActivity As Activity = Nothing
            Me.clonedInstanceStateStream.Position = 0
            [Using] New RuntimeEnvironment(Me.workflowExecutor.WorkflowRuntime)
                rootActivity = Activity.Load(Me.clonedInstanceStateStream, Me.workflowDefinition)
            End [Using]
            rootActivity.SetValue(WorkflowExecutor.TrackingListenerBrokerProperty, Me.workflowExecutor.RootActivity.GetValue(WorkflowExecutor.TrackingListenerBrokerProperty))
            Dim newWorkflowExecutor As New WorkflowExecutor(Guid.Empty)
            newWorkflowExecutor.Initialize(rootActivity, Me.workflowExecutor.WorkflowRuntime, Me.workflowExecutor)
            Dim contextActivityForId As Activity = newWorkflowExecutor.GetContextActivityForId(Me.activityContextId)
            Dim activityByName As Activity = contextActivityForId.GetActivityByName(Me.activityQualifiedName)
            [Using] New ServiceEnvironment(activityByName)
                [Using] newWorkflowExecutor.SetCurrentActivity(activityByName)
                    [Using] context As ActivityExecutionContext = New ActivityExecutionContext(activityByName)
                        context.Invoke([Of] EventArgs)(Me.callbackHandler, Me.callbackData)
                    End [Using]
                End [Using]
            End [Using]
            newWorkflowExecutor.BatchCollection.WorkItemOrderId = Me.workflowExecutor.BatchCollection.WorkItemOrderId
            Dim pair As KeyValuePair([Of] Object, WorkBatch)
            For Each pair In Me.workflowExecutor.BatchCollection
                pair.Value.SetWorkBatchCollection(newWorkflowExecutor.BatchCollection)
                Dim key As Activity = [TryCast](pair.Key,Activity)
                If (Not key Is Nothing) Then
                    Dim activity5 As Activity = contextActivityForId.GetActivityByName(key.QualifiedName)
                    newWorkflowExecutor.BatchCollection.Add(activity5, pair.Value)
                End If
            Next
            Me.workflowExecutor.BatchCollection.Clear
            newWorkflowExecutor.CompletedContextActivities = Me.completedContextActivities
            Me.workflowExecutor.WorkflowRuntime.ReplaceWorkflowExecutor(Me.workflowExecutor.InstanceId, Me.workflowExecutor, newWorkflowExecutor)
            If Not Me.suspendOnRevert Then
                newWorkflowExecutor.Scheduler.Resume
            Else
                newWorkflowExecutor.SuspendOnIdle(Me.suspendOnRevertInfo)
            End If
            Me.DisposeCheckpointState
        End Sub


        ' Properties
        Friend ReadOnly Property IsInstanceStateRevertRequested As Boolean
            Get
                Return Me.isInstanceStateRevertRequested
            End Get
        End Property


        ' Fields
        Private activityContextId As Integer
        Private activityQualifiedName As String
        Private callbackData As EventArgs
        Private callbackHandler As EventHandler([Of] EventArgs)
        Private clonedInstanceStateStream As MemoryStream
        Private completedContextActivities As Hashtable = New Hashtable
        Private isInstanceStateRevertRequested As Boolean
        Private suspendOnRevert As Boolean
        Private suspendOnRevertInfo As String
        Private workflowDefinition As Activity
        Private workflowExecutor As WorkflowExecutor
    End Class
End Namespace

